#include "CNN/cifar10.hpp"

#include <torch/data/example.h>
#include <torch/types.h>

#include <fstream>

constexpr const uint32_t BATCH_NUM = 5;
constexpr const uint32_t BATCH_ITEM_NUM = 10000;
constexpr const uint32_t BATCH_IDX_POS = 11;
constexpr const char *kTrainBaseFilename = "data_batch_1.bin";
constexpr const char *kTestFilename = "test_batch.bin";

constexpr uint32_t kTrainSize = BATCH_NUM * BATCH_ITEM_NUM;
constexpr uint32_t kTestSize = 10000;

constexpr uint32_t kImageChannels = 3;
constexpr uint32_t kImageRows = 32;
constexpr uint32_t kImageColumns = 32;
constexpr uint32_t kImageSize = kImageChannels * kImageRows * kImageColumns;

constexpr double eps = 1e-5;

std::string join_paths(std::string head, const std::string &tail) {
  if (head.back() != '/') {
    head.push_back('/');
  }
  head += tail;
  return head;
}

void read_train_items(torch::Tensor &images, torch::Tensor &labels,
                      const std::string &root, bool check) {
  images = torch::empty({kTrainSize, kImageChannels, kImageRows, kImageColumns},
                        torch::kByte);
  labels = torch::empty(kTrainSize, torch::kByte);

  double RGB[3] = {0.0, 0.0, 0.0};
  for (uint32_t i = 0; i < BATCH_NUM; i++) {
    std::string batch_file = kTrainBaseFilename;
    batch_file[BATCH_IDX_POS] += i;
    const auto path = join_paths(root, batch_file);
    std::ifstream items(path, std::ios::binary);
    TORCH_CHECK(items, "Error opening images file at ", path);

    const uint32_t offset = i * BATCH_ITEM_NUM;
    for (uint32_t j = 0; j < BATCH_ITEM_NUM; j++) {
      items.read(reinterpret_cast<char *>(labels.data_ptr()) + offset + j, 1);
      items.read(reinterpret_cast<char *>(images.data_ptr()) +
                     (offset + j) * kImageSize,
                 kImageSize);
      if (check) {
        auto p = reinterpret_cast<uint8_t *>(images.data_ptr()) +
                 (offset + j) * kImageSize;
        for (uint32_t c = 0; c < kImageChannels; c++) {
          for (uint32_t k = 0; k < kImageRows * kImageColumns; k++) {
            RGB[c] += p[c * kImageRows * kImageColumns + k] / 255.0;
          }
        }
      }
    }
  }

  if (check) {
    for (uint32_t c = 0; c < kImageChannels; c++) {
      RGB[c] /= BATCH_NUM * BATCH_ITEM_NUM * kImageRows * kImageColumns;
      assert(std::fabs(RGB[c] - CIFAR10::CIFAR10_MEANS[c]) < eps &&
             "Error means");
    }
  }

  images = std::move(images.to(torch::kFloat32));
  images -= images.min();
  images /= images.max();

  labels = std::move(labels.to(torch::kInt64));
}

void read_test_items(torch::Tensor &images, torch::Tensor &labels,
                     const std::string &root) {
  images = torch::empty({kTestSize, kImageChannels, kImageRows, kImageColumns},
                        torch::kByte);
  labels = torch::empty(kTestSize, torch::kByte);

  const auto path = join_paths(root, kTestFilename);
  std::ifstream items(path, std::ios::binary);
  TORCH_CHECK(items, "Error opening images file at ", path);
  for (uint32_t i = 0; i < BATCH_ITEM_NUM; i++) {
    items.read(reinterpret_cast<char *>(labels.data_ptr()) + i, 1);
    items.read(reinterpret_cast<char *>(images.data_ptr()) + i * kImageSize,
               kImageSize);
  }

  images = std::move(images.to(torch::kFloat32));
  images -= images.min();
  images /= images.max();

  labels = std::move(labels.to(torch::kInt64));
}

CIFAR10::CIFAR10(const std::string &root, Mode mode, bool check) {
  if (mode == Mode::kTrain) {
    read_train_items(images_, targets_, root, check);
  } else {
    read_test_items(images_, targets_, root);
  }
}

torch::data::Example<> CIFAR10::get(size_t index) {
  return {images_[index], targets_[index]};
}

torch::optional<size_t> CIFAR10::size() const { return images_.size(0); }

bool CIFAR10::is_train() const noexcept {
  return images_.size(0) == kTrainSize;
}

const torch::Tensor &CIFAR10::images() const { return images_; }

const torch::Tensor &CIFAR10::targets() const { return targets_; }